package com.flink.hbase.statemanagement.config;

import org.apache.flink.configuration.Configuration;
import org.apache.flink.configuration.CheckpointingOptions;
import org.apache.flink.configuration.StateBackendOptions;
import org.apache.flink.contrib.streaming.state.RocksDBStateBackend;
import org.apache.flink.contrib.streaming.state.RocksDBOptions;
import org.apache.flink.runtime.state.StateBackend;
import org.apache.flink.runtime.state.filesystem.FsStateBackend;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
import org.apache.flink.streaming.api.CheckpointingMode;
import org.apache.flink.streaming.api.environment.CheckpointConfig;
import org.apache.flink.runtime.state.memory.MemoryStateBackend;
import org.rocksdb.CompactionStyle;
import org.rocksdb.InfoLogLevel;

import java.io.IOException;
import java.time.Duration;

/**
 * RocksDB状态后端配置类
 * 
 * 提供高性能的状态管理配置，包括：
 * 1. RocksDB状态后端配置
 * 2. Checkpoint配置
 * 3. 性能优化配置
 * 4. 不同环境的配置预设
 * 
 * @author Flink Team
 */
public class StateBackendConfig {
    
    // 默认配置常量
    private static final String DEFAULT_CHECKPOINT_DIR = "hdfs://namenode:9000/flink/checkpoints";
    private static final String DEFAULT_SAVEPOINT_DIR = "hdfs://namenode:9000/flink/savepoints";
    private static final long DEFAULT_CHECKPOINT_INTERVAL = 30000L; // 30秒
    private static final long DEFAULT_CHECKPOINT_TIMEOUT = 600000L; // 10分钟
    private static final long DEFAULT_MIN_PAUSE_BETWEEN_CHECKPOINTS = 5000L; // 5秒
    
    /**
     * 配置基础的RocksDB状态后端
     * 
     * @param env Flink执行环境
     * @param checkpointDir 检查点目录
     * @return 配置后的环境
     */
    public static StreamExecutionEnvironment configureBasicRocksDB(
            StreamExecutionEnvironment env, String checkpointDir) {
        
        try {
            // 创建RocksDB状态后端
            RocksDBStateBackend rocksDBStateBackend = new RocksDBStateBackend(checkpointDir);
            
            // 基础配置
            rocksDBStateBackend.setDbStoragePath("/tmp/rocksdb");
            rocksDBStateBackend.setNumberOfTransferThreads(4);
            rocksDBStateBackend.setNumberOfTransferingThreads(4);
            
            // 设置状态后端
            env.setStateBackend(rocksDBStateBackend);
            
            // 启用检查点
            env.enableCheckpointing(DEFAULT_CHECKPOINT_INTERVAL);
            
            // 配置检查点
            CheckpointConfig checkpointConfig = env.getCheckpointConfig();
            checkpointConfig.setCheckpointingMode(CheckpointingMode.EXACTLY_ONCE);
            checkpointConfig.setCheckpointTimeout(DEFAULT_CHECKPOINT_TIMEOUT);
            checkpointConfig.setMinPauseBetweenCheckpoints(DEFAULT_MIN_PAUSE_BETWEEN_CHECKPOINTS);
            checkpointConfig.setMaxConcurrentCheckpoints(1);
            checkpointConfig.enableExternalizedCheckpoints(
                CheckpointConfig.ExternalizedCheckpointCleanup.RETAIN_ON_CANCELLATION);
            
            System.out.println("✅ 基础RocksDB状态后端配置完成");
            return env;
            
        } catch (Exception e) {
            System.err.println("❌ RocksDB状态后端配置失败: " + e.getMessage());
            throw new RuntimeException("配置RocksDB状态后端失败", e);
        }
    }
    
    /**
     * 配置高性能RocksDB状态后端
     * 
     * @param env Flink执行环境
     * @param checkpointDir 检查点目录
     * @return 配置后的环境
     */
    public static StreamExecutionEnvironment configureHighPerformanceRocksDB(
            StreamExecutionEnvironment env, String checkpointDir) {
        
        try {
            // 创建配置对象
            Configuration config = new Configuration();
            
            // 状态后端配置
            config.setString(StateBackendOptions.STATE_BACKEND, "rocksdb");
            config.setString(CheckpointingOptions.CHECKPOINTS_DIRECTORY, checkpointDir);
            config.setString(CheckpointingOptions.SAVEPOINT_DIRECTORY, DEFAULT_SAVEPOINT_DIR);
            
            // RocksDB性能优化配置
            config.setString(RocksDBOptions.LOCAL_DIRECTORIES, "/tmp/rocksdb,/data/rocksdb");
            config.setLong(RocksDBOptions.BLOCK_CACHE_SIZE, 256 * 1024 * 1024L); // 256MB
            config.setLong(RocksDBOptions.WRITE_BUFFER_SIZE, 64 * 1024 * 1024L); // 64MB
            config.setInteger(RocksDBOptions.MAX_WRITE_BUFFER_NUMBER, 4);
            config.setInteger(RocksDBOptions.MIN_WRITE_BUFFER_NUMBER_TO_MERGE, 2);
            config.setInteger(RocksDBOptions.MAX_BACKGROUND_JOBS, 4);
            config.setInteger(RocksDBOptions.MAX_BACKGROUND_COMPACTIONS, 2);
            config.setInteger(RocksDBOptions.MAX_BACKGROUND_FLUSHES, 2);
            config.setInteger(RocksDBOptions.TARGET_FILE_SIZE_BASE, 64 * 1024 * 1024); // 64MB
            config.setInteger(RocksDBOptions.MAX_BYTES_FOR_LEVEL_BASE, 256 * 1024 * 1024); // 256MB
            config.setString(RocksDBOptions.COMPACTION_STYLE, CompactionStyle.LEVEL.name());
            config.setString(RocksDBOptions.LOG_LEVEL, InfoLogLevel.WARN_LEVEL.name());
            config.setBoolean(RocksDBOptions.USE_BLOOM_FILTER, true);
            config.setLong(RocksDBOptions.BLOOM_FILTER_BITS_PER_KEY, 10);
            config.setLong(RocksDBOptions.BLOOM_FILTER_BLOCK_BASED_MODE, true);
            
            // 应用配置
            env.configure(config);
            
            // 创建优化的RocksDB状态后端
            RocksDBStateBackend rocksDBStateBackend = new RocksDBStateBackend(checkpointDir);
            rocksDBStateBackend.setDbStoragePath("/tmp/rocksdb");
            rocksDBStateBackend.setNumberOfTransferThreads(8);
            rocksDBStateBackend.setNumberOfTransferingThreads(8);
            
            env.setStateBackend(rocksDBStateBackend);
            
            // 高性能检查点配置
            env.enableCheckpointing(10000L); // 10秒检查点间隔
            
            CheckpointConfig checkpointConfig = env.getCheckpointConfig();
            checkpointConfig.setCheckpointingMode(CheckpointingMode.EXACTLY_ONCE);
            checkpointConfig.setCheckpointTimeout(300000L); // 5分钟超时
            checkpointConfig.setMinPauseBetweenCheckpoints(2000L); // 2秒最小间隔
            checkpointConfig.setMaxConcurrentCheckpoints(2);
            checkpointConfig.setTolerableCheckpointFailureNumber(3);
            checkpointConfig.enableExternalizedCheckpoints(
                CheckpointConfig.ExternalizedCheckpointCleanup.RETAIN_ON_CANCELLATION);
            checkpointConfig.enableUnalignedCheckpoints(true); // 启用非对齐检查点
            
            System.out.println("🚀 高性能RocksDB状态后端配置完成");
            return env;
            
        } catch (Exception e) {
            System.err.println("❌ 高性能RocksDB状态后端配置失败: " + e.getMessage());
            throw new RuntimeException("配置高性能RocksDB状态后端失败", e);
        }
    }
    
    /**
     * 配置开发环境的RocksDB状态后端
     * 
     * @param env Flink执行环境
     * @return 配置后的环境
     */
    public static StreamExecutionEnvironment configureDevRocksDB(StreamExecutionEnvironment env) {
        try {
            // 开发环境使用本地文件系统
            String checkpointDir = "file:///tmp/flink-checkpoints";
            
            RocksDBStateBackend rocksDBStateBackend = new RocksDBStateBackend(checkpointDir);
            rocksDBStateBackend.setDbStoragePath("/tmp/rocksdb-dev");
            
            env.setStateBackend(rocksDBStateBackend);
            
            // 开发环境检查点配置
            env.enableCheckpointing(60000L); // 1分钟检查点间隔
            
            CheckpointConfig checkpointConfig = env.getCheckpointConfig();
            checkpointConfig.setCheckpointingMode(CheckpointingMode.EXACTLY_ONCE);
            checkpointConfig.setCheckpointTimeout(120000L); // 2分钟超时
            checkpointConfig.setMinPauseBetweenCheckpoints(10000L); // 10秒最小间隔
            checkpointConfig.setMaxConcurrentCheckpoints(1);
            checkpointConfig.enableExternalizedCheckpoints(
                CheckpointConfig.ExternalizedCheckpointCleanup.DELETE_ON_CANCELLATION);
            
            System.out.println("🛠️ 开发环境RocksDB状态后端配置完成");
            return env;
            
        } catch (Exception e) {
            System.err.println("❌ 开发环境RocksDB状态后端配置失败: " + e.getMessage());
            throw new RuntimeException("配置开发环境RocksDB状态后端失败", e);
        }
    }
    
    /**
     * 配置生产环境的RocksDB状态后端
     * 
     * @param env Flink执行环境
     * @param checkpointDir 检查点目录
     * @param savepointDir 保存点目录
     * @return 配置后的环境
     */
    public static StreamExecutionEnvironment configureProductionRocksDB(
            StreamExecutionEnvironment env, String checkpointDir, String savepointDir) {
        
        try {
            // 生产环境配置
            Configuration config = new Configuration();
            
            // 状态后端配置
            config.setString(StateBackendOptions.STATE_BACKEND, "rocksdb");
            config.setString(CheckpointingOptions.CHECKPOINTS_DIRECTORY, checkpointDir);
            config.setString(CheckpointingOptions.SAVEPOINT_DIRECTORY, savepointDir);
            
            // 生产环境RocksDB优化配置
            config.setString(RocksDBOptions.LOCAL_DIRECTORIES, 
                "/data1/rocksdb,/data2/rocksdb,/data3/rocksdb");
            config.setLong(RocksDBOptions.BLOCK_CACHE_SIZE, 512 * 1024 * 1024L); // 512MB
            config.setLong(RocksDBOptions.WRITE_BUFFER_SIZE, 128 * 1024 * 1024L); // 128MB
            config.setInteger(RocksDBOptions.MAX_WRITE_BUFFER_NUMBER, 6);
            config.setInteger(RocksDBOptions.MIN_WRITE_BUFFER_NUMBER_TO_MERGE, 3);
            config.setInteger(RocksDBOptions.MAX_BACKGROUND_JOBS, 8);
            config.setInteger(RocksDBOptions.MAX_BACKGROUND_COMPACTIONS, 4);
            config.setInteger(RocksDBOptions.MAX_BACKGROUND_FLUSHES, 4);
            config.setInteger(RocksDBOptions.TARGET_FILE_SIZE_BASE, 128 * 1024 * 1024); // 128MB
            config.setInteger(RocksDBOptions.MAX_BYTES_FOR_LEVEL_BASE, 512 * 1024 * 1024); // 512MB
            config.setString(RocksDBOptions.COMPACTION_STYLE, CompactionStyle.LEVEL.name());
            config.setString(RocksDBOptions.LOG_LEVEL, InfoLogLevel.INFO_LEVEL.name());
            config.setBoolean(RocksDBOptions.USE_BLOOM_FILTER, true);
            config.setLong(RocksDBOptions.BLOOM_FILTER_BITS_PER_KEY, 10);
            config.setLong(RocksDBOptions.BLOOM_FILTER_BLOCK_BASED_MODE, true);
            
            // 应用配置
            env.configure(config);
            
            // 创建生产环境RocksDB状态后端
            RocksDBStateBackend rocksDBStateBackend = new RocksDBStateBackend(checkpointDir);
            rocksDBStateBackend.setDbStoragePath("/data/rocksdb");
            rocksDBStateBackend.setNumberOfTransferThreads(16);
            rocksDBStateBackend.setNumberOfTransferingThreads(16);
            
            env.setStateBackend(rocksDBStateBackend);
            
            // 生产环境检查点配置
            env.enableCheckpointing(5000L); // 5秒检查点间隔
            
            CheckpointConfig checkpointConfig = env.getCheckpointConfig();
            checkpointConfig.setCheckpointingMode(CheckpointingMode.EXACTLY_ONCE);
            checkpointConfig.setCheckpointTimeout(600000L); // 10分钟超时
            checkpointConfig.setMinPauseBetweenCheckpoints(1000L); // 1秒最小间隔
            checkpointConfig.setMaxConcurrentCheckpoints(3);
            checkpointConfig.setTolerableCheckpointFailureNumber(5);
            checkpointConfig.enableExternalizedCheckpoints(
                CheckpointConfig.ExternalizedCheckpointCleanup.RETAIN_ON_CANCELLATION);
            checkpointConfig.enableUnalignedCheckpoints(true); // 启用非对齐检查点
            
            System.out.println("🏭 生产环境RocksDB状态后端配置完成");
            return env;
            
        } catch (Exception e) {
            System.err.println("❌ 生产环境RocksDB状态后端配置失败: " + e.getMessage());
            throw new RuntimeException("配置生产环境RocksDB状态后端失败", e);
        }
    }
    
    /**
     * 配置内存状态后端（用于对比）
     * 
     * @param env Flink执行环境
     * @return 配置后的环境
     */
    public static StreamExecutionEnvironment configureMemoryStateBackend(
            StreamExecutionEnvironment env) {
        
        try {
            // 创建内存状态后端
            MemoryStateBackend memoryStateBackend = new MemoryStateBackend();
            memoryStateBackend.setMaxStateSize(100 * 1024 * 1024); // 100MB
            
            env.setStateBackend(memoryStateBackend);
            
            // 内存状态后端检查点配置
            env.enableCheckpointing(30000L);
            
            CheckpointConfig checkpointConfig = env.getCheckpointConfig();
            checkpointConfig.setCheckpointingMode(CheckpointingMode.EXACTLY_ONCE);
            checkpointConfig.setCheckpointTimeout(60000L);
            checkpointConfig.setMinPauseBetweenCheckpoints(5000L);
            checkpointConfig.setMaxConcurrentCheckpoints(1);
            
            System.out.println("💾 内存状态后端配置完成");
            return env;
            
        } catch (Exception e) {
            System.err.println("❌ 内存状态后端配置失败: " + e.getMessage());
            throw new RuntimeException("配置内存状态后端失败", e);
        }
    }
    
    /**
     * 配置文件系统状态后端（用于对比）
     * 
     * @param env Flink执行环境
     * @param checkpointDir 检查点目录
     * @return 配置后的环境
     */
    public static StreamExecutionEnvironment configureFsStateBackend(
            StreamExecutionEnvironment env, String checkpointDir) {
        
        try {
            // 创建文件系统状态后端
            FsStateBackend fsStateBackend = new FsStateBackend(checkpointDir);
            fsStateBackend.setCheckpointStreamFactory(null); // 使用默认
            
            env.setStateBackend(fsStateBackend);
            
            // 文件系统状态后端检查点配置
            env.enableCheckpointing(20000L);
            
            CheckpointConfig checkpointConfig = env.getCheckpointConfig();
            checkpointConfig.setCheckpointingMode(CheckpointingMode.EXACTLY_ONCE);
            checkpointConfig.setCheckpointTimeout(300000L);
            checkpointConfig.setMinPauseBetweenCheckpoints(3000L);
            checkpointConfig.setMaxConcurrentCheckpoints(2);
            
            System.out.println("📁 文件系统状态后端配置完成");
            return env;
            
        } catch (Exception e) {
            System.err.println("❌ 文件系统状态后端配置失败: " + e.getMessage());
            throw new RuntimeException("配置文件系统状态后端失败", e);
        }
    }
    
    /**
     * 打印状态后端配置信息
     * 
     * @param env Flink执行环境
     */
    public static void printStateBackendInfo(StreamExecutionEnvironment env) {
        StateBackend stateBackend = env.getStateBackend();
        System.out.println("\n=== 状态后端配置信息 ===");
        System.out.println("状态后端类型: " + stateBackend.getClass().getSimpleName());
        
        CheckpointConfig checkpointConfig = env.getCheckpointConfig();
        System.out.println("检查点间隔: " + checkpointConfig.getCheckpointInterval() + "ms");
        System.out.println("检查点模式: " + checkpointConfig.getCheckpointingMode());
        System.out.println("检查点超时: " + checkpointConfig.getCheckpointTimeout() + "ms");
        System.out.println("最大并发检查点: " + checkpointConfig.getMaxConcurrentCheckpoints());
        System.out.println("最小间隔: " + checkpointConfig.getMinPauseBetweenCheckpoints() + "ms");
        System.out.println("外部检查点: " + checkpointConfig.isExternalizedCheckpointsEnabled());
        System.out.println("非对齐检查点: " + checkpointConfig.isUnalignedCheckpointsEnabled());
        System.out.println("========================\n");
    }
    
    /**
     * 配置环境特定的状态后端
     * 
     * @param env Flink执行环境
     * @param environment 环境类型：dev, test, prod
     * @param checkpointDir 检查点目录
     * @return 配置后的环境
     */
    public static StreamExecutionEnvironment configureByEnvironment(
            StreamExecutionEnvironment env, String environment, String checkpointDir) {
        
        switch (environment.toLowerCase()) {
            case "dev":
                return configureDevRocksDB(env);
            case "test":
                return configureBasicRocksDB(env, checkpointDir);
            case "prod":
                return configureProductionRocksDB(env, checkpointDir, DEFAULT_SAVEPOINT_DIR);
            default:
                System.out.println("⚠️ 未知环境类型: " + environment + "，使用默认配置");
                return configureBasicRocksDB(env, checkpointDir);
        }
    }
} 